import os
import numpy as np
import re

# 预先计算的权重矩阵加载
weights = {
    'left_upper': np.load('E:\\wafer52\\weight\\Top_Left.npy'),
    'right_upper': np.load('E:\\wafer52\\weight\\Top_Right.npy'),
    'left_lower': np.load('E:\\wafer52\\weight\\Bottom_Left.npy'),
    'right_lower': np.load('E:\\wafer52\\weight\\Bottom_Right.npy')
}


def get_sub_flow_files(row, col):
    return [
        (row, col, 'left_upper'),
        (row, col + 1, 'right_upper'),
        (row + 1, col, 'left_lower'),
        (row + 1, col + 1, 'right_lower')
    ]


def merge_flows(flow_dir, output_dir, tile_size=256, full_size=512):
    os.makedirs(output_dir, exist_ok=True)
    flow_files = [f for f in os.listdir(flow_dir) if f.endswith('.npy')]

    for flow_file in flow_files:
        match = re.search(r'tile_\((\d+)_(\d+),(\d+)_(\d+)\)\.npy', flow_file)
        if match:
            row1, col1, row2, col2 = map(int, match.groups())
            parent_path = os.path.join(flow_dir, flow_file)
            flow = np.load(parent_path)

            sub_flows = get_sub_flow_files(row1, col1)
            for (row, col, weight_key) in sub_flows:
                weight_matrix = weights[weight_key][:, :, np.newaxis]
                sub_flow = flow[:tile_size, :tile_size] * weight_matrix

                output_file = os.path.join(output_dir, f"tile_({row}_{col}).npy")
                if os.path.exists(output_file):
                    existing_flow = np.load(output_file)
                    sub_flow += existing_flow

                np.save(output_file, sub_flow)
                print(f"Processed {output_file}")


# 示例用法
flow_dir = 'E:\\wafer52\\11867_32nm_fine_align_512_dvf(1)'  # 替换为实际输入目录
output_dir = 'E:\\wafer52\\11867_32nm_fine_align_256_dvf_repair(2)'  # 替换为实际输出目录

merge_flows(flow_dir, output_dir, tile_size=256, full_size=512)
